import cv2
import xml.etree.ElementTree as ET
import random
import os
import math
import pandas as pd


# 坐标变换函数
def fun(x, y):
    # 定义矩形的四个点坐标
    dt = {}
    x1, y1 = 0, 0
    x2, y2 = x, 0
    x3, y3 = x, y
    x4, y4 = 0, y
    for angle in range(1, 90):
        # 计算矩形的中心点坐标
        cx = (x1 + x2 + x3 + x4) / 4
        cy = (y1 + y2 + y3 + y4) / 4

        # 将点移到以中心点为原点的坐标系中
        x1 -= cx
        y1 -= cy
        x2 -= cx
        y2 -= cy
        x3 -= cx
        y3 -= cy
        x4 -= cx
        y4 -= cy

        # 将角度转换为弧度
        angle_rad = math.radians(angle)

        # 使用旋转矩阵进行旋转
        x1_new = x1 * math.cos(angle_rad) - y1 * math.sin(angle_rad)
        y1_new = x1 * math.sin(angle_rad) + y1 * math.cos(angle_rad)
        x2_new = x2 * math.cos(angle_rad) - y2 * math.sin(angle_rad)
        y2_new = x2 * math.sin(angle_rad) + y2 * math.cos(angle_rad)
        x3_new = x3 * math.cos(angle_rad) - y3 * math.sin(angle_rad)
        y3_new = x3 * math.sin(angle_rad) + y3 * math.cos(angle_rad)
        x4_new = x4 * math.cos(angle_rad) - y4 * math.sin(angle_rad)
        y4_new = x4 * math.sin(angle_rad) + y4 * math.cos(angle_rad)

        # 返回到原坐标系中
        x1_new += cx
        y1_new += cy
        x2_new += cx
        y2_new += cy
        x3_new += cx
        y3_new += cy
        x4_new += cx
        y4_new += cy

        # # 打印旋转后的矩形四个点坐标
        # print("旋转后的矩形四个点坐标：")
        # print(f"点1: ({x1_new}, {y1_new})")
        # print(f"点2: ({x2_new}, {y2_new})")
        # print(f"点3: ({x3_new}, {y3_new})")
        # print(f"点4: ({x4_new}, {y4_new})")
        # list1.append((abs(x4_new-x2_new), abs(y3_new-y1_new)))#返回方框长度
        y = [y1_new, y2_new, y3_new, y4_new]
        x = [x1_new, x2_new, x3_new, x4_new]
        xx = max(x) - min(x)
        yy = max(y) - min(y)
        k = abs(yy) / abs(xx)
        dt[k] = (xx, yy)

    return dt


battery_model = fun(200, 124)
can_model = fun(544, 229)
bottle_model = fun(550, 200)


# 1. 解析 VOC 标签
def parse_voc_label(xml_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()

    objects = []
    for obj in root.findall('object'):
        name = obj.find('name').text
        bbox = obj.find('bndbox')
        xmin = int(bbox.find('xmin').text)
        ymin = int(bbox.find('ymin').text)
        xmax = int(bbox.find('xmax').text)
        ymax = int(bbox.find('ymax').text)
        objects.append({'name': name, 'bbox': (xmin, ymin, xmax, ymax)})

    return objects


profiles = "input/"
out_profiles = "output/"

imgs = os.listdir(profiles + 'img/')
for img in imgs:
    print(img)
    image = cv2.imread(profiles + r'img/' + img)
    try:
        labels = parse_voc_label(profiles + 'ano/' + img.split('.')[0] + '.xml')
    except:
        continue
    flag = 0
    for label in labels:
        background = cv2.imread("background.png")
        # 4. 创建新图像区域
        ojb_name = label['name']
        xmin, ymin, xmax, ymax = label['bbox']
        object_region1 = image[ymin:ymax, xmin:xmax]

        # 4_2 等比压缩图像
        if ojb_name == 'battery':
            k1 = abs((ymax - ymin) / (xmax - xmin))  # 得到横纵长度比例
            loss = pd.Series(None)
            for bm_k, value in battery_model.items():
                loss.loc[abs(bm_k - k1)] = value
            loss.sort_index(inplace=True)
            xy = loss.iloc[0]  # 获得比例最接近的方框的坐标
            k = (xy[0] + 5 - random.randint(0, 10)) / (xmax - xmin)
            if k > 2 or k < 0.5:
                print(k1)
                print(k)
                print(xy)
            object_region = cv2.resize(object_region1, (0, 0), fx=k, fy=k)  # 按比例缩放
            del loss

        if ojb_name == 'can':
            k1 = abs((ymax - ymin) / (xmax - xmin))
            loss = pd.Series(None)
            for bm_k, value in can_model.items():
                loss.loc[abs(bm_k - k1)] = value
            loss.sort_index(inplace=True)
            xy = loss.iloc[0]  # 获得比例最接近的方框的坐标
            k = (xy[0] + 5 - random.randint(0, 10)) / (xmax - xmin)
            if k > 2 or k < 0.5:
                print(k1)
                print(k)
                print(xy)
            object_region = cv2.resize(object_region1, (0, 0), fx=k, fy=k)  # 按比例缩放
            del loss

        if ojb_name == 'bottle':
            k1 = abs((ymax - ymin) / (xmax - xmin))
            loss = pd.Series(None)
            for bm_k, value in bottle_model.items():
                loss.loc[abs(bm_k - k1)] = value
            loss.sort_index(inplace=True)
            xy = loss.iloc[0]  # 获得比例最接近的方框的坐标
            k = (xy[0] + 5 - random.randint(0, 10)) / (xmax - xmin)
            if k > 2 or k < 0.5:
                print(k1)
                print(k)
                print(xy)
            object_region = cv2.resize(object_region1, (0, 0), fx=k, fy=k)  # 按比例缩放
            del loss

        # 5. 在新背景上创建图像
        object_height, object_width, _ = object_region.shape
        try:
            new_xmin = random.randint(258, 1003 - object_width)
            new_ymin = random.randint(70, 914 - object_height)
        except:
            print('尺寸过大:', img)
            continue
        new_xmax = new_xmin + object_width
        new_ymax = new_ymin + object_height
        new_bg = background[new_ymin:new_ymin + object_height, new_xmin:new_xmin + object_width]
        # print(new_ymin, new_ymin + object_height, new_xmin, new_xmin + object_width)
        # 6. 复制物体到新图像上
        for i in range(object_height):
            for j in range(object_width):
                if object_region[i, j].any():
                    new_bg[i, j] = object_region[i, j]
        label['bbox'] = (new_xmin, new_ymin, new_xmax, new_ymax)

        # 9. 创建新的 XML 文件
        root = ET.Element('annotation')
        obj = ET.SubElement(root, 'object')
        name = ET.SubElement(obj, 'name')
        name.text = label['name']
        bndbox = ET.SubElement(obj, 'bndbox')
        xmin, ymin, xmax, ymax = label['bbox']
        ET.SubElement(bndbox, 'xmin').text = str(xmin)
        ET.SubElement(bndbox, 'ymin').text = str(ymin)
        ET.SubElement(bndbox, 'xmax').text = str(xmax)
        ET.SubElement(bndbox, 'ymax').text = str(ymax)

        # 10. 保存新的 XML 文件
        tree = ET.ElementTree(root)
        tree.write(out_profiles + 'ano/' + img.split('.')[0] + '_' + str(flag) + '.xml')
        # 7. 保存结果图像
        cv2.imwrite(out_profiles + 'img/' + img.split('.')[0] + '_' + str(flag) + '.jpg', background)
        flag += 1
